# tools/mcp_client.py
# MCP客户端管理器，负责与MCP服务器的通信

from typing import List, Dict, Any, Optional


class MCPClientManager:
    """MCP客户端管理器，维护长期连接"""
    
    def __init__(self):
        self.session = None
        self.client = None
        self.read = None
        self.write = None
        
    async def initialize(self, server_path: str) -> bool:
        """初始化MCP客户端连接（兼容旧版本）"""
        return await self.initialize_with_command("python", [server_path], {})
    
    async def initialize_with_command(self, command: str, args: List[str], env: Dict[str, str] = None) -> bool:
        """使用完整命令初始化MCP客户端连接 - 改进错误处理"""
        try:
            from mcp.client.stdio import stdio_client
            from mcp.client.session import ClientSession
            from mcp import StdioServerParameters
            from mcp.shared.exceptions import McpError
            
            # 创建服务器参数
            server_params = StdioServerParameters(
                command=command,
                args=args,
                env=env
            )
            
            # 创建MCP客户端连接
            self.client = stdio_client(server_params)
            self.read, self.write = await self.client.__aenter__()
            self.session = ClientSession(self.read, self.write)
            await self.session.__aenter__()
            
            # 尝试初始化，捕获更详细的错误
            try:
                await self.session.initialize()
                return True
            except McpError as e:
                # MCP协议错误，通常是服务器启动失败或连接关闭
                error_msg = f"MCP协议错误: {e}"
                if "Connection closed" in str(e):
                    error_msg = f"连接已关闭，可能是服务器启动失败或命令错误: {command} {' '.join(args)}"
                print(f"MCP客户端初始化失败: {error_msg}")
                return False
            except Exception as e:
                error_msg = f"初始化过程中发生错误: {e}"
                print(f"MCP客户端初始化失败: {error_msg}")
                import traceback
                traceback.print_exc()
                return False
            
        except FileNotFoundError:
            error_msg = f"命令未找到: {command}，请确保已安装并添加到PATH"
            print(f"MCP客户端初始化失败: {error_msg}")
            return False
        except Exception as e:
            error_msg = f"MCP客户端初始化失败: {e}"
            print(f"{error_msg}")
            import traceback
            traceback.print_exc()
            return False
    
    async def call_tool(self, tool_name: str, kwargs: Dict[str, Any]) -> str:
        """调用MCP工具"""
        try:
            if not self.session:
                raise RuntimeError("MCP会话未初始化")
            
            print(f"[MCP调试] 调用工具: {tool_name}, 参数: {kwargs}")
            
            # 修复参数传递：MCP工具期望的是arguments格式
            result = await self.session.call_tool(tool_name, arguments=kwargs)
            print(f"[MCP调试] 工具返回: {result}")
            
            if result.content and len(result.content) > 0:
                return result.content[0].text
            else:
                return "工具执行成功，但无返回内容"
        except Exception as e:
            print(f"[MCP调试] 工具调用错误: {e}")
            return f"工具调用失败: {str(e)}"
    
    async def list_tools(self) -> List[Any]:
        """获取可用工具列表"""
        if not self.session:
            return []
        try:
            tools_result = await self.session.list_tools()
            return tools_result.tools
        except Exception as e:
            print(f"获取工具列表失败: {e}")
            return []
    
    async def close(self):
        """关闭MCP连接 - 深度修复异步错误"""
        import asyncio
        
        try:
            # 先关闭session
            if self.session:
                try:
                    await asyncio.wait_for(
                        self.session.__aexit__(None, None, None),
                        timeout=2.0
                    )
                except (asyncio.TimeoutError, asyncio.CancelledError):
                    # 忽略超时和取消错误
                    pass
                except Exception:
                    # 忽略其他异常
                    pass
            
            # 再关闭client
            if self.client:
                try:
                    if hasattr(self.client, '__aexit__'):
                        await asyncio.wait_for(
                            self.client.__aexit__(None, None, None),
                            timeout=2.0
                        )
                except (asyncio.TimeoutError, asyncio.CancelledError):
                    # 忽略超时和取消错误
                    pass
                except Exception:
                    # 忽略其他异常
                    pass
        except Exception:
            # 最终异常保护
            pass
        finally:
            # 确保清理引用
            self.session = None
            self.client = None
            self.read = None
            self.write = None
